import os
import logging
from typing import List, Dict, Any, Optional, Callable
from pathlib import Path
from langchain.text_splitter import (
    RecursiveCharacterTextSplitter,
    CharacterTextSplitter,
    TokenTextSplitter,
    MarkdownHeaderTextSplitter
)
from langchain.schema import Document
from .document_processor import DocumentProcessor
from .embedding_manager import EmbeddingManager
from config import Config
from config.thresholds import KnowledgeBaseThresholds, RAGThresholds, ValidationThresholds

logger = logging.getLogger(__name__)

class KnowledgeBaseConfig:
    """知识库配置类"""
    
    def __init__(
        self,
        name: str,
        description: str = "",
        chunk_size: int = 1000,
        chunk_overlap: int = 200,
        splitter_type: str = "recursive",
        custom_separators: List[str] = None,
        metadata_filters: Dict[str, Any] = None
    ):
        self.name = name
        self.description = description
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap
        self.splitter_type = splitter_type
        self.custom_separators = custom_separators or ["\n\n", "\n", "。", "！", "？", ".", "!", "?", " ", ""]
        self.metadata_filters = metadata_filters or {}

class KnowledgeBaseManager:
    """知识库管理器"""
    
    def __init__(self, vector_db_path: str = None):
        self.vector_db_path = vector_db_path or Config.VECTOR_DB_PATH
        self.document_processor = DocumentProcessor(vector_db_path)
        self.embedding_manager = EmbeddingManager()
        self.knowledge_bases: Dict[str, KnowledgeBaseConfig] = {}
        
        # 确保目录存在
        os.makedirs(self.vector_db_path, exist_ok=True)
        
        # 自动发现和加载现有的知识库
        self._discover_existing_knowledge_bases()
    
    def create_knowledge_base(
        self,
        name: str,
        description: str = "",
        chunk_size: int = KnowledgeBaseThresholds.DEFAULT_CHUNK_SIZE,
        chunk_overlap: int = KnowledgeBaseThresholds.DEFAULT_CHUNK_OVERLAP,
        splitter_type: str = "recursive",
        custom_separators: List[str] = None,
        metadata_filters: Dict[str, Any] = None
    ) -> Dict[str, Any]:
        """创建新的知识库"""
        try:
            logger.info(f"创建知识库: {name}")
            
            # 检查知识库是否已存在
            if name in self.knowledge_bases:
                return {
                    "status": "success",
                    "message": f"知识库 '{name}' 已存在，无需重复创建",
                    "knowledge_base": {
                        "name": name,
                        "description": description,
                        "chunk_size": chunk_size,
                        "chunk_overlap": chunk_overlap,
                        "splitter_type": splitter_type,
                        "document_count": 0
                    }
                }
            
            # 创建知识库配置
            kb_config = KnowledgeBaseConfig(
                name=name,
                description=description,
                chunk_size=chunk_size,
                chunk_overlap=chunk_overlap,
                splitter_type=splitter_type,
                custom_separators=custom_separators,
                metadata_filters=metadata_filters
            )
            
            # 保存配置
            self.knowledge_bases[name] = kb_config
            
            # 创建空的向量存储
            self._create_empty_collection(name)
            
            result = {
                "status": "success",
                "message": f"知识库 '{name}' 创建成功",
                "knowledge_base": {
                    "name": name,
                    "description": description,
                    "chunk_size": chunk_size,
                    "chunk_overlap": chunk_overlap,
                    "splitter_type": splitter_type,
                    "document_count": 0
                }
            }
            
            logger.info(f"知识库创建成功: {name}")
            return result
            
        except Exception as e:
            logger.error(f"创建知识库失败: {str(e)}")
            return {
                "status": "error",
                "message": f"创建知识库失败: {str(e)}"
            }
    
    def upload_document_to_kb(
        self,
        file_path: str,
        knowledge_base_name: str,
        custom_metadata: Dict[str, Any] = None,
        auto_create: bool = True,
        kb_config: Dict[str, Any] = None
    ) -> Dict[str, Any]:
        """上传文档到指定知识库"""
        try:
            logger.info(f"上传文档到知识库: {knowledge_base_name}")
            
            # 检查知识库是否存在
            if knowledge_base_name not in self.knowledge_bases:
                if auto_create:
                    logger.info(f"知识库 '{knowledge_base_name}' 不存在，自动创建")
                    
                    # 使用提供的配置或默认配置创建知识库
                    create_config = kb_config or {
                        "description": f"自动创建的知识库: {knowledge_base_name}",
                        "chunk_size": KnowledgeBaseThresholds.DEFAULT_CHUNK_SIZE,
                        "chunk_overlap": KnowledgeBaseThresholds.DEFAULT_CHUNK_OVERLAP,
                        "splitter_type": "recursive"
                    }
                    
                    create_result = self.create_knowledge_base(
                        name=knowledge_base_name,
                        description=create_config.get("description", ""),
                        chunk_size=create_config.get("chunk_size", KnowledgeBaseThresholds.DEFAULT_CHUNK_SIZE),
                        chunk_overlap=create_config.get("chunk_overlap", KnowledgeBaseThresholds.DEFAULT_CHUNK_OVERLAP),
                        splitter_type=create_config.get("splitter_type", "recursive"),
                        custom_separators=create_config.get("custom_separators"),
                        metadata_filters=create_config.get("metadata_filters")
                    )
                    
                    if create_result["status"] != "success":
                        return create_result
                    
                    logger.info(f"知识库 '{knowledge_base_name}' 创建成功")
                else:
                    return {
                        "status": "error",
                        "message": f"知识库 '{knowledge_base_name}' 不存在，请先创建知识库"
                    }
            
            kb_config = self.knowledge_bases[knowledge_base_name]
            
            # 使用知识库特定的配置处理文档
            result = self._process_document_with_config(
                file_path, 
                kb_config, 
                custom_metadata
            )
            
            return result
            
        except Exception as e:
            logger.error(f"上传文档失败: {str(e)}")
            return {
                "status": "error",
                "message": f"上传文档失败: {str(e)}"
            }
    
    def _process_document_with_config(
        self,
        file_path: str,
        kb_config: KnowledgeBaseConfig,
        custom_metadata: Dict[str, Any] = None
    ) -> Dict[str, Any]:
        """使用特定配置处理文档"""
        try:
            # 1. 加载文档
            documents = self.document_processor.load_document(file_path)
            
            # 2. 添加自定义元数据
            if custom_metadata:
                for doc in documents:
                    doc.metadata.update(custom_metadata)
                    doc.metadata["knowledge_base"] = kb_config.name
            
            # 3. 使用知识库特定的拆分器
            split_docs = self._split_documents_with_config(documents, kb_config)
            
            # 4. 存储到向量数据库
            try:
                vector_store = self.document_processor.add_documents_to_store(
                    split_docs, 
                    kb_config.name
                )
            except:
                # 如果集合不存在，创建新的
                vector_store = self.document_processor.create_vector_store(
                    split_docs, 
                    kb_config.name
                )
            
            # 5. 获取集合信息
            collection_info = self.document_processor.get_collection_info(kb_config.name)
            
            return {
                "status": "success",
                "message": f"文档上传到知识库 '{kb_config.name}' 成功",
                "knowledge_base": kb_config.name,
                "collection_info": collection_info,
                "file_path": file_path,
                "chunks_created": len(split_docs)
            }
            
        except Exception as e:
            logger.error(f"处理文档失败: {str(e)}")
            raise
    
    def _split_documents_with_config(
        self,
        documents: List[Document],
        kb_config: KnowledgeBaseConfig
    ) -> List[Document]:
        """使用特定配置拆分文档"""
        try:
            logger.info(f"使用配置拆分文档: {kb_config.name}")
            
            # 根据配置选择拆分器
            if kb_config.splitter_type == "recursive":
                splitter = RecursiveCharacterTextSplitter(
                    chunk_size=kb_config.chunk_size,
                    chunk_overlap=kb_config.chunk_overlap,
                    separators=kb_config.custom_separators
                )
            elif kb_config.splitter_type == "character":
                splitter = CharacterTextSplitter(
                    chunk_size=kb_config.chunk_size,
                    chunk_overlap=kb_config.chunk_overlap,
                    separator="\n"
                )
            elif kb_config.splitter_type == "token":
                splitter = TokenTextSplitter(
                    chunk_size=kb_config.chunk_size,
                    chunk_overlap=kb_config.chunk_overlap
                )
            elif kb_config.splitter_type == "markdown":
                # 先按Markdown头部拆分，再按字符拆分
                headers_to_split_on = [
                    ("#", "标题1"),
                    ("##", "标题2"),
                    ("###", "标题3"),
                ]
                markdown_splitter = MarkdownHeaderTextSplitter(
                    headers_to_split_on=headers_to_split_on
                )
                char_splitter = CharacterTextSplitter(
                    chunk_size=kb_config.chunk_size,
                    chunk_overlap=kb_config.chunk_overlap
                )
                
                # 先按Markdown头部拆分
                markdown_splits = []
                for doc in documents:
                    splits = markdown_splitter.split_text(doc.page_content)
                    for split in splits:
                        new_doc = Document(
                            page_content=split.page_content,
                            metadata=doc.metadata.copy()
                        )
                        new_doc.metadata.update(split.metadata)
                        markdown_splits.append(new_doc)
                
                # 再按字符拆分
                split_docs = char_splitter.split_documents(markdown_splits)
            else:
                # 默认使用递归拆分
                splitter = RecursiveCharacterTextSplitter(
                    chunk_size=kb_config.chunk_size,
                    chunk_overlap=kb_config.chunk_overlap,
                    separators=kb_config.custom_separators
                )
                split_docs = splitter.split_documents(documents)
            
            if kb_config.splitter_type != "markdown":
                split_docs = splitter.split_documents(documents)
            
            logger.info(f"文档拆分完成，生成 {len(split_docs)} 个文本块")
            return split_docs
            
        except Exception as e:
            logger.error(f"拆分文档失败: {str(e)}")
            raise
    
    def ask_question_in_kb(
        self,
        question: str,
        knowledge_base_name: str,
        top_k: int = None,
        rerank_threshold: float = RAGThresholds.DEFAULT_RERANK_THRESHOLD,
        user_id: str = None,
        session_id: str = None
    ) -> Dict[str, Any]:
        """在指定知识库中回答问题"""
        try:
            # 检查知识库是否存在
            if knowledge_base_name not in self.knowledge_bases:
                return {
                    "status": "error",
                    "message": f"知识库 '{knowledge_base_name}' 不存在"
                }
            
            # 使用RAG系统回答问题
            from .rag_system import RAGSystem
            rag_system = RAGSystem(self.vector_db_path)
            
            result = rag_system.ask_question(
                question=question,
                collection_name=knowledge_base_name,
                top_k=top_k,
                rerank_threshold=rerank_threshold,
                user_id=user_id,
                session_id=session_id
            )
            
            # 添加知识库信息
            if result["status"] == "success":
                result["knowledge_base"] = knowledge_base_name
                result["kb_config"] = {
                    "name": self.knowledge_bases[knowledge_base_name].name,
                    "description": self.knowledge_bases[knowledge_base_name].description
                }
            
            return result
            
        except Exception as e:
            logger.error(f"在知识库中回答问题失败: {str(e)}")
            return {
                "status": "error",
                "message": f"回答问题失败: {str(e)}",
                "question": question
            }
    
    def list_knowledge_bases(self) -> Dict[str, Any]:
        """列出所有知识库"""
        try:
            kb_list = []
            for name, config in self.knowledge_bases.items():
                # 获取每个知识库的文档数量
                try:
                    collection_info = self.document_processor.get_collection_info(name)
                    doc_count = collection_info.get("document_count", 0)
                except:
                    doc_count = 0
                
                kb_info = {
                    "name": name,
                    "description": config.description,
                    "chunk_size": config.chunk_size,
                    "chunk_overlap": config.chunk_overlap,
                    "splitter_type": config.splitter_type,
                    "document_count": doc_count
                }
                kb_list.append(kb_info)
            
            return {
                "status": "success",
                "knowledge_bases": kb_list,
                "count": len(kb_list)
            }
            
        except Exception as e:
            logger.error(f"列出知识库失败: {str(e)}")
            return {
                "status": "error",
                "message": f"列出知识库失败: {str(e)}"
            }
    
    def delete_knowledge_base(self, name: str) -> Dict[str, Any]:
        """删除知识库"""
        try:
            if name not in self.knowledge_bases:
                return {
                    "status": "error",
                    "message": f"知识库 '{name}' 不存在"
                }
            
            # 删除知识库配置
            del self.knowledge_bases[name]
            
            # 删除向量存储（这里需要根据具体的向量数据库实现）
            # 暂时只删除配置
            
            return {
                "status": "success",
                "message": f"知识库 '{name}' 删除成功"
            }
            
        except Exception as e:
            logger.error(f"删除知识库失败: {str(e)}")
            return {
                "status": "error",
                "message": f"删除知识库失败: {str(e)}"
            }
    
    def _discover_existing_knowledge_bases(self):
        """自动发现和加载现有的知识库"""
        try:
            logger.info("开始自动发现现有知识库...")
            
            # 获取所有集合
            collections = self.document_processor.list_collections()
            
            for collection_name in collections:
                if collection_name not in self.knowledge_bases:
                    logger.info(f"发现现有知识库: {collection_name}")
                    
                    # 为现有集合创建默认配置
                    kb_config = KnowledgeBaseConfig(
                        name=collection_name,
                        description=f"自动发现的现有知识库: {collection_name}",
                        chunk_size=KnowledgeBaseThresholds.DEFAULT_CHUNK_SIZE,
                        chunk_overlap=KnowledgeBaseThresholds.DEFAULT_CHUNK_OVERLAP,
                        splitter_type="recursive"
                    )
                    
                    # 添加到知识库管理器
                    self.knowledge_bases[collection_name] = kb_config
                    
                    logger.info(f"成功加载现有知识库: {collection_name}")
            
            logger.info(f"自动发现完成，共发现 {len(collections)} 个知识库")
            
        except Exception as e:
            logger.error(f"自动发现知识库失败: {str(e)}")
    
    def _create_empty_collection(self, collection_name: str):
        """创建空的集合"""
        try:
            # 这里可以创建空的向量存储集合
            # 具体实现取决于向量数据库
            pass
        except Exception as e:
            logger.warning(f"创建空集合失败: {str(e)}") 